#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
数据准备主脚本

用于运行数据准备流水线，将原始数据转换为机器学习模型可用的格式
"""

import os
import sys
import argparse
import json
from datetime import datetime
from src.data.data_preparation_controller import DataPreparationController


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="股票数据准备工具")
    
    parser.add_argument("--config", type=str, default="config/data_preparation_config.json",
                       help="配置文件路径，默认为config/data_preparation_config.json")
    parser.add_argument("--stocks", type=str, nargs="+",
                       help="股票代码列表，如果不指定则使用配置文件中的默认值")
    parser.add_argument("--start_date", type=str,
                       help="开始日期，格式为YYYY-MM-DD，如果不指定则使用配置文件中的默认值")
    parser.add_argument("--end_date", type=str,
                       help="结束日期，格式为YYYY-MM-DD，如果不指定则使用配置文件中的默认值")
    parser.add_argument("--output", type=str,
                       help="输出目录路径，如果不指定则使用配置文件中的默认值")
    parser.add_argument("--create_config", action="store_true",
                       help="创建默认配置文件")
    
    args = parser.parse_args()
    
    # 创建必要的目录
    os.makedirs("logs", exist_ok=True)
    os.makedirs("data/prepared", exist_ok=True)
    os.makedirs("config", exist_ok=True)
    
    # 处理创建配置文件的请求
    if args.create_config:
        create_default_config(args.config)
        return
    
    # 检查配置文件是否存在
    if not os.path.exists(args.config):
        print(f"配置文件 {args.config} 不存在，将创建默认配置文件")
        create_default_config(args.config)
    
    # 创建数据准备控制器
    controller = DataPreparationController(args.config)
    
    # 如果指定了输出目录，更新配置
    if args.output:
        controller.config['prepared_data_path'] = args.output
        os.makedirs(args.output, exist_ok=True)
    
    # 运行数据准备流水线
    print(f"开始数据准备流程: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    result = controller.prepare_data_pipeline(
        stock_codes=args.stocks,
        start_date=args.start_date,
        end_date=args.end_date
    )
    print(f"数据准备完成: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    
    # 输出结果
    print(f"\n数据已保存至: {result['data_path']}")
    print(f"数据集大小:")
    print(f"  训练集: {result['shapes']['train']}")
    print(f"  验证集: {result['shapes']['valid']}")
    print(f"  测试集: {result['shapes']['test']}")
    
    # 检查是否有数据
    if result['shapes']['train'][0] == 0:
        print("\n警告: 训练集为空，请检查数据源和参数配置")


def create_default_config(config_path):
    """创建默认配置文件"""
    default_config = {
        "price_api_key": "",
        "volume_api_key": "",
        "sentiment_model_path": "",
        "db_connection_string": "",
        "default_stock_codes": ["AAPL", "MSFT", "GOOG", "AMZN", "TSLA"],
        "default_start_date": "2020-01-01",
        "default_end_date": "2023-12-31",
        "train_ratio": 0.7,
        "valid_ratio": 0.15,
        "test_ratio": 0.15,
        "sequence_length": 10,
        "prepared_data_path": "./data/prepared",
        "feature_engineering": {
            "normalization_method": "standard",
            "n_selected_features": 20,
            "lag_periods": [1, 3, 5, 7],
            "prediction_horizon": 5,
            "price_change_threshold": 0.02
        },
        "logging": {
            "level": "INFO",
            "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
            "log_file": "logs/data_preparation.log"
        }
    }
    
    # 保存到文件
    with open(config_path, 'w', encoding='utf-8') as f:
        json.dump(default_config, f, indent=4)
    
    print(f"已创建默认配置文件: {config_path}")
    print("请根据需要修改配置文件中的参数，特别是API密钥和数据库连接字符串")


if __name__ == "__main__":
    main() 